package com.zhny.test;

import org.apache.parquet.Strings;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.mllib.clustering.KMeans;
import org.apache.spark.mllib.clustering.KMeansModel;
import org.apache.spark.mllib.linalg.Vector;
import org.apache.spark.mllib.linalg.Vectors;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
//k均值
public class KMeansUtil {
    public static void exc(String dataFilePath, String resultFilePath) {
        SparkConf sparkConf = new SparkConf().setAppName("NaiveBayesTest").setMaster("local[1]");
        sparkConf.set("spark.driver.allowMultipleContexts", "true");
        JavaSparkContext sc = new JavaSparkContext(sparkConf);

        JavaRDD<String> data = null;

        if (Strings.isNullOrEmpty(dataFilePath)) {
            data = sc.textFile("src/main/resources/data/KMeansData.txt");
        } else {
            data = sc.textFile(dataFilePath);
        }

        JavaRDD<Vector> parsedData = data.map(new Function<String, Vector>() {
            private static final long serialVersionUID = 1L;
            @Override
            public Vector call(String line) throws Exception {
                String[] dsStr = line.split(" ");

                double[] ds = new double[dsStr.length];

                for (int i = 0; i < dsStr.length; i++) {
                    ds[i] = Double.parseDouble(dsStr[i]);
                }

                return Vectors.dense(ds);
            }
        }).cache();

        int numIterations = 100;

        KMeansModel kMeansModel = KMeans.train(parsedData.rdd(), 2, numIterations, 10);

        FileWriter fos = null;
        try {
            fos = new FileWriter(new File(resultFilePath));
            fos.write(kMeansModel.toPMML() + "\n");

            double[] d = new double[] { 78, 12, 66, 25 };
            double[] d2 = new double[] { 24, 78, 15, 66 };
            double[] d3 = new double[] { 14, 56, 37, 88 };
            double[] d4 = new double[] { 65, 22, 76, 35 };
            Vector v = Vectors.dense(d);
            Vector v2 = Vectors.dense(d2);
            Vector v3 = Vectors.dense(d3);
            Vector v4 = Vectors.dense(d4);

            fos.write(String.valueOf(new Double(kMeansModel.predict(v)).longValue()) + "\n");
            fos.write(String.valueOf(new Double(kMeansModel.predict(v2)).longValue()) + "\n");
            fos.write(String.valueOf(new Double(kMeansModel.predict(v3)).longValue()) + "\n");
            fos.write(String.valueOf(new Double(kMeansModel.predict(v4)).longValue()) + "\n");

            fos.flush();
            fos.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
